# -*-coding=utf-8-*-
"""
结合Gaussian程序的Dimer
"""

import numpy as np
from pes import PES
import os
import re
import random
import time


class GaussianFile():
    """
    gaussian文件处理
    """
    def __init__(self):
        self.pattern_force = re.compile(r'\s+\d+\s+\d+\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)\s+')
        self.pattern_spe = re.compile(r'SCF Done:  E\(.+?\) =\s+(-?\d+?\.\d*E?-?\d*)')
        self.patter_coordinates = re.compile(r' (\w+)\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)\s+')
        self.file_head = ['%mem=500MB\n',
                          '%nprocshared=4\n',
                          '#p force pm6\n',
                          '\n',
                          'title\n',
                          '\n',
                          '0 1\n']

    def init_read(self, file_path):
        """
        初始化读取
        """
        if os.path.exists(file_path):
            with open(file_path, 'r') as f:
                file_lines = f.readlines()
        else:
            print('文件不存在! :', file_path)
            return
        elements = []
        coordinates = []
        for line in file_lines:
            result = re.search(self.patter_coordinates, line)
            if result:
                elements.append(result.group(1))
                coordinates.append([float(result.group(2)),
                                    float(result.group(3)),
                                    float(result.group(4))])
            elif len(elements) > 0:
                return (elements, np.array(coordinates, dtype=float))

    def read_file(self, file_path):
        """
        读取文件并解析
        """
        if os.path.exists(file_path):
            with open(file_path, 'r') as f:
                file_lines = f.readlines()
        else:
            print('文件不存在! :', file_path)
            return
        current_line = 0
        # search SPE
        spe = None
        for i, line in enumerate(file_lines):
            result = re.search(self.pattern_spe, line)
            if result:
                spe = float(result.group(1))
        # search force
        for i, line in enumerate(file_lines):
            if line[:42] == ' ***** Axes restored to original set *****':
                current_line = i + 5
        force_data = []
        for i in range(current_line, len(file_lines)):
            result = re.search(self.pattern_force, file_lines[i])
            if result:
                # collecting force data
                force_data.append([float(result.group(1)),
                                   float(result.group(2)),
                                   float(result.group(3))])
            else:
                # end
                return (spe, force_data)

    def generate_input(self, path, coordinates, elements):
        """
        生成输入文件
        """
        with open(path, 'w') as f:
            for line in self.file_head:
                f.writelines(line)
            for i in range(len(coordinates)):
                f.writelines(' '+elements[i]+'  '+str(coordinates[i][0])+'  '+str(coordinates[i][1])+'  '+str(coordinates[i][2])+'\n')
            f.writelines('\n')

    def run_gaussian(self, path, input_file):
        """
        计算
        """
        os.chdir(path)
        print('running gaussian: ', input_file)
        os.system('g09 '+input_file)
        print('done')


class Dimer():
    def __init__(self, n, elements, position, path, vector=None):
        self.min_value = 1e-16
        self.path = path
        self.gaussian = GaussianFile()
        self.n = n
        self.position = np.array(position, dtype='f')
        self.elements = elements
        # dimer半径
        self.r = 0.0001
        if vector is None:
            self.vector = np.zeros(self.n, 'f')
            self.vector[:] = np.random.rand(self.n)
            self.vector[:] -= 0.5
        else:
            self.vector = np.array(vector, 'f')
        # 向量归一化
        # 这里用=/会报错
        self.vector = self.vector / np.linalg.norm(self.vector)
        self.normal = np.zeros(self.n, 'f')
        self.f1 = np.zeros(self.n, 'f')
        self.f2 = np.zeros(self.n, 'f')
        # 曲率
        self.c = 0
        self.update_c()
        self.delta = 0.001
        self.delta_angle = np.pi / 180
        # dimer速度
        self.v = np.zeros(self.n, 'f')
        # dimer步进时间
        self.timer = 0.3
        self.angle = 0
        # 计数器
        self.step_number = 0
        self.rotate_number = 0

    def get_value(self, position):
        """
        获取值
        """
        return self.PES.get_value(position)

    def force(self, position, name):
        """
        计算像点受力
        """
        path = os.path.join(self.path, name)
        # 一维转二维数组
        new_position = position.reshape(int(len(position)/3), 3)
        self.gaussian.generate_input(path, new_position, self.elements)
        self.gaussian.run_gaussian(self.path, name)
        e, force = self.gaussian.read_file(path[:-3]+'out')
        force = np.array(force, 'f')
        # 二维转一维
        force = force.flatten()
        return np.array(force, 'f')
        # return self.PES.get_force(position)

    def vertical_force(self, force, vector):
        """
        计算垂直分力
        参数：
            force：力向量
            vector：dimer向量
        """
        # 平行分力
        f_parallel = np.dot(vector, force) * vector
        # 垂直分力
        f_vertical = force - f_parallel
        return f_vertical

    def get_rotate_angle(self):
        """
        计算dimer的旋转角
        """
        self.f1[:] = self.force(self.position+self.vector*self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_1.gjf')
        self.f2[:] = self.force(self.position-self.vector*self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_2.gjf')
        print(self.f1-self.f2)
        print(np.dot(np.dot(self.f1-self.f2, self.vector), self.vector))
        delta_f = (self.f1-self.f2) - np.dot(np.dot(self.f1-self.f2, self.vector), self.vector)
        # 力大小
        f_abs = np.linalg.norm(delta_f)
        if f_abs < self.min_value:
            return 0
        # 垂直向量判断
        self.update_normal()
        if np.isnan(self.normal[0]):
            return 0
        # 这里是单位向量
        new_vector = self.vector*np.cos(self.delta_angle) + self.normal*np.sin(self.delta_angle)
        # 微旋转后的dimer受力
        new_f1 = self.force(self.position+new_vector * self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_1_.gjf')
        new_f2 = self.force(self.position-new_vector * self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_2_.gjf')
        delta_new_f = (new_f1-new_f2) - np.dot(np.dot(new_f1-new_f2, new_vector), new_vector)
        temp = f_abs*np.cos(2*self.delta_angle)-np.dot(delta_new_f, delta_f) / f_abs
        angle = np.arctan((np.sin(2*self.delta_angle)*f_abs) / temp) / 2
        return angle

    def rotate(self, angle):
        """
        旋转dimer
        """
        self.angle = angle
        # 垂直向量
        if np.isnan(self.normal[0]):
            return
        new_vector = self.vector*np.cos(self.angle) + self.normal*np.sin(self.angle)
        self.vector[:] = new_vector / np.linalg.norm(new_vector)
        self.f1 = self.force(self.position+self.vector*self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_1_.gjf')
        self.f2 = self.force(self.position-self.vector*self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_2_.gjf')
        # 计算曲率
        self.update_c()

    def translate(self):
        """
        移动dimer
        """
        self.angle = 0
        # dimer合力
        f_r = (self.f1 + self.f2) / 2
        # dimer平行力
        f_parallel = np.dot(self.vector, f_r) * self.vector
        # dimer指向鞍点力
        if self.c < 0:
            f_to_saddle = f_r - f_parallel * 2
        else:
            f_to_saddle = - f_parallel
        delta_v = f_to_saddle * self.timer
        # 速度调整
        if np.dot(self.v, f_to_saddle) < 0:
            self.v = delta_v
        else:
            delta_v_abs = np.dot(delta_v, delta_v)
            if delta_v_abs < self.min_value:
                self.v[:] = 0
            else:
                self.v = delta_v * (1+np.dot(delta_v, self.v)/np.dot(delta_v, delta_v))
        self.position += (self.v * self.timer)
        # 计算新点相关数据
        self.f1[:] = self.force(self.position+self.vector*self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_1_.gjf')
        self.f2[:] = self.force(self.position-self.vector*self.r, str(self.step_number)+'_'+str(self.rotate_number)+'_2_.gjf')

    def work(self):
        # 旋转到曲率最小
        # rotate_angle = np.pi / 180
        # self.position[:] = [np.random.random(), np.random.random()]
        x = [self.position[0]]
        y = [self.position[1]]
        times = []
        for i in range(120):
            self.rotate_number = 0
            for j in range(20):
                rotate_angle = self.get_rotate_angle()
                self.rotate(rotate_angle)
                self.rotate_number += 1
                # print(rotate_angle)
                # print(np.linalg.norm(self.vertical_force(self.f1-self.f2, self.vector)))
                # 垂直力大小
                if np.linalg.norm(self.vertical_force(self.f1-self.f2, self.vector)) < 0.001:
                    times.append(j)
                    break
                elif j == 99:
                    times.append(j)
            # print(self.vector)
            self.translate()
            x.append(self.position[0])
            y.append(self.position[1])
            if (np.abs(self.f1 + self.f2) < 0.1).all():
                print('step: ', i)
                break
            self.step_number += 1
        print('time', times)
        print(self.position/np.pi*180)

    def update_c(self):
        """
        计算曲率
        """
        self.c = -np.dot(self.f1-self.f2, self.vector) / 2 / self.r

    def update_normal(self):
        """
        计算法向量
        """
        normal = self.vertical_force(self.f1-self.f2, self.vector)
        normal_abs = np.linalg.norm(normal)
        if normal_abs < self.min_value:
            # 如果垂直力为0，中断计算垂直向量
            self.normal[:] = np.float('nan')
        else:
            self.normal[:] = normal / normal_abs


if __name__ == "__main__":
    gauss = GaussianFile()
    path = r'C:\Users\Administrator\Desktop\dimertest'
    init_file = 'init.gjf'
    init_file2 = 'init2.gjf'
    e1, cord1 = gauss.init_read(os.path.join(path, init_file))
    e2, cord2 = gauss.init_read(os.path.join(path, init_file2))
    cord1 = cord1.flatten()
    cord2 = cord2.flatten()
    vector = cord1 - cord2
    d = Dimer(len(e1)*3, e1, cord1, path, vector)
    d.work()
